package main

import (
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"golang.org/x/crypto/blowfish"
	"io"
	"log"
	"os"
)

func blowfishEncrypt(key []byte) func([]byte) []byte {
	cipher, err := blowfish.NewCipher(key)
	if err != nil {
		log.Fatalf("failed to create cipher: %v", err)
	}
	return func(data []byte) []byte {
		if len(data) != 8 {
			log.Fatalf("invalid block size: %d", len(data))
		}
		// Endian swap before encryption
		block := make([]byte, 8)
		binary.BigEndian.PutUint32(block[0:4], binary.LittleEndian.Uint32(data[0:4]))
		binary.BigEndian.PutUint32(block[4:8], binary.LittleEndian.Uint32(data[4:8]))

		cipher.Encrypt(block, block)

		// Endian swap after encryption
		result := make([]byte, 8)
		binary.LittleEndian.PutUint32(result[0:4], binary.BigEndian.Uint32(block[0:4]))
		binary.LittleEndian.PutUint32(result[4:8], binary.BigEndian.Uint32(block[4:8]))

		return result
	}
}
func decrypt(f io.Reader, pw string, encoding string) (string, error) {
	// Validate the file header
	header := make([]byte, 12)
	_, err := f.Read(header)
	if err != nil {
		return "", fmt.Errorf("failed to read header: %v", err)
	}
	if string(header) != "VimCrypt~03!" {
		return "", fmt.Errorf("not a blowfish2-encoded vimcrypt file")
	}

	// Read salt
	salt := make([]byte, 8)
	_, err = f.Read(salt)
	if err != nil {
		return "", fmt.Errorf("failed to read salt: %v", err)
	}

	// Derive the key
	for i := 0; i < 1000; i++ {
		pw = fmt.Sprintf("%x", sha256.Sum256(append([]byte(pw), salt...)))
	}
	fmt.Println(pw)

	// Encrypt the password using SHA256 and salt
	key := sha256.Sum256(append([]byte(pw), salt...))
	//// 打印生成的 key 值
	//fmt.Printf("Generated key (hex): %s\n", hex.EncodeToString(key[:])) // 打印 key
	cipher := blowfishEncrypt(key[:]) // 使用 key[:] 转换为切片

	// Read the initial blocks
	block0 := make([]byte, 8)
	//block1 := make([]byte, 8)
	_, err = f.Read(block0)
	if err != nil {
		return "", fmt.Errorf("failed to read block0: %v", err)
	}
	//_, err = f.Read(block1)
	//if err != nil {
	//	return "", fmt.Errorf("failed to read block1: %v", err)
	//}

	// Decrypt the file content
	var decrypted []byte
	for {
		block1 := make([]byte, 8)
		n, err := f.Read(block1) // n 是实际读到的字节数

		// block1 is now n bytes, pad or trim
		if n == 0 && err == io.EOF {
			break
		}
		if err != nil && err != io.EOF {
			return "", fmt.Errorf("failed to read block1: %v", err)
		}

		//fmt.Printf("Inputed block(hex): %s\n", hex.EncodeToString(block0)) // 打印 key
		encryptedBlock := cipher(block0)
		//fmt.Printf("Generated block(hex): %s\n", hex.EncodeToString(encryptedBlock)) // 打印 key
		//fmt.Printf("block1(hex): %s\n", hex.EncodeToString(block1))                  // 打印 key
		for i := 0; i < n; i++ {
			//for i := 0; i < len(block1); i++ {
			decrypted = append(decrypted, encryptedBlock[i]^block1[i])
		}
		//block0 = block1
		copy(block0, block1)
		//copy(block1, make([]byte, 8))
		//_, err = f.Read(block1)
		if err != nil {
			if err == io.EOF {
				break
			}
			return "", fmt.Errorf("failed to read block1: %v", err)
		}
	}
	//fmt.Printf("Decrypted string(hex): %s\n", hex.EncodeToString(decrypted)) // 打印 key
	return string(decrypted), nil
}

func main() {
	// Test the decryption function
	filePath := "test.txt" // Path to the encrypted file
	pw := "Gao123"

	file, err := os.Open(filePath)
	if err != nil {
		log.Fatalf("failed to open file: %v", err)
	}
	defer file.Close()

	result, err := decrypt(file, pw, "utf8")
	if err != nil {
		log.Fatalf("decryption failed: %v", err)
	}

	fmt.Println("Decrypted text:", result)
}
